import caffe
import surgery, score

import numpy as np
import os
import sys
import logging

FORMAT = '[%(levelname)-5s]%(asctime)-8s %(filename)s:%(lineno)d %(message)s'
DATEFORMAT = '%H:%M:%S'
logging.basicConfig(level=logging.DEBUG, format=FORMAT, datefmt=DATEFORMAT)
logging.debug('start')


try:
    import setproctitle
    setproctitle.setproctitle(os.path.basename(os.getcwd()))
except:
    pass

weights = 'fcn-alexnet-pascal.caffemodel'

# init
caffe.set_device(int(sys.argv[1]))
caffe.set_mode_gpu()

solver = caffe.SGDSolver('solver.prototxt')
if len(sys.argv) > 2:
    solver.restore(sys.argv[2])
else:
    solver.net.copy_from(weights)

# surgeries
interp_layers = [k for k in solver.net.params.keys() if 'up' in k]
surgery.interp(solver.net, interp_layers)

# scoring
# val = np.loadtxt('../data/segvalid11.txt', dtype=str)

for _ in range(25000):
    solver.step(100)
    # score.seg_tests(solver, False, val, layer='score')
